// The MPC formulation for r^2 as in Ni Dong's dissertation.
#include "MPC_EW2D_M2_v2_FiniteHorizon.h"
using namespace Ipopt;
// -----------------------------------------------------------------------
/* Constructor. */
MPC_Formulation::MPC_Formulation()
{
	int i,m,n,i_mode,j_mode;
	r2_set = 2.25;
	//q_r = 1/(r2_set*r2_set);
	q_r = 0.0;
	m2_set = 600;
	q_m = 1/(m2_set*m2_set);
	
	Tl = 400.0;
	Th = 900.0;
	dTl = -50.0;
	dTh = 50.0;

	t_dep = 200.0;
	t_k   = 0.0;
	dt = 5.0;
	P = 5;

	MPC_Formulation::mode = 20;
	k_max = 100;
	numUniqueCOV = mode*(mode)-1;

	T_old = 610;
	W     = 1.0;


	z = new double[mode*mode];
	Km2 = new double*[mode];
	double dx = pi/k_max;
	double temp = dx*dx*pi*pi;
	for(m=0;m<mode;m++)
	{
		Km2[m] = new double[mode];
		for(n=0;n<mode;n++)
		{
			Km2[m][n] = (2-2*cos(2*m*dx))/temp;
			z[index2D(m,n)] = 0.0;
		}
	}

	iCon_r_begin = 0;
	iCon_m_begin = P;
	iCon_c2_begin = 2*P;
	iCon_sigma2_begin = 3*P;
	iCon_cov_begin = 4*P;
	iCon_dT_begin = 4*P+P*(mode*mode-1);

	i_T_begin = 0;
	i_r_begin = P;
	i_m_begin = 2*P;
	i_c2_begin = 3*P;
	i_sigma2_begin = 4*P;
	i_cov_begin = 5*P;

	iJacG_r_r_begin = 0;
	iJacG_r_cov_begin = P;
	iJacG_m_m_begin = P+P*numUniqueCOV;
	iJacG_m_cov_begin = 2*P+P*numUniqueCOV;
	iJacG_c_T_begin = 2*P+2*P*numUniqueCOV;
	iJacG_c_c_begin = 3*P+2*P*numUniqueCOV;
	iJacG_sigma_T_begin = 4*P+2*P*numUniqueCOV;
	iJacG_sigma_sigma_begin = 5*P+2*P*numUniqueCOV;
	iJacG_cov_c_begin = 6*P+2*P*numUniqueCOV;
	iJacG_cov_sigma_begin = 6*P+3*P*numUniqueCOV;
	iJacG_cov_cov_cur_begin = 6*P+4*P*numUniqueCOV;
	iJacG_cov_cov_pre_begin = 6*P+4*P*numUniqueCOV+P*numUniqueCOV;
	iJacG_dT_T_cur_begin = 6*P+4*P*numUniqueCOV+(2*P-1)*numUniqueCOV;
	iJacG_dT_T_pre_begin = 6*P+4*P*numUniqueCOV+(2*P-1)*numUniqueCOV+P;

	CONST_COEFF_C2 = k_c/(k_max*k_max)*pow(W,-a_c);
	CONST_COEFF_SIGMA2 = pi*pi/(k_max*k_max);

	x_old = new double[(5+mode*mode-1)*P];

	// we initialize x in bounds, in the upper right quadrant
	for(i=0;i<P;i++)
	{
		x_old[i_T(i)] = T_old;
		x_old[i_r(i)] = 0.0;
		x_old[i_m(i)] = 0.0;
		for(i_mode=0;i_mode<mode;i_mode++){
			for(j_mode=0;j_mode<mode;j_mode++){
				if(i_mode==j_mode && i_mode == 0){}
				else{
					x_old[i_r(i)] += z[index2D(i_mode,j_mode)];
					x_old[i_m(i)] += Km2[i_mode][j_mode]*z[index2D(i_mode,j_mode)];
					x_old[i_cov(i,i_mode,j_mode)] = 0.0;
				}
			}
		}
		x_old[i_r(i)] /= (pi*pi);
		x_old[i_c2(i)] = CONST_COEFF_C2*exp(k_B*x_old[i_T(i)]/E_c);
		x_old[i_sigma2(i)] = CONST_COEFF_SIGMA2*(1+exp((a_t+k_t*W)*x_old[i_T(i)]-(a_v+k_v*W)));
	}
}
// -----------------------------------------------------------------------
MPC_Formulation::~MPC_Formulation()
{
	delete []x_old;
	delete []z;
	int m;
	for(m=0;m<mode;m++){
		delete [] Km2[m];
	}
	delete [] Km2;
}
// -----------------------------------------------------------------------
void MPC_Formulation::parameterize(double t, const double* z_cov)
{
	int m,n;
	double COV_mn;
	for(m=0;m<mode;m++){
		for(n=0;n<mode;n++){
			COV_mn = z_cov[index2D(m,n)];
			z[index2D(m,n)] = z_cov[index2D(m,n)];
		}
	}
	MPC_Formulation::time = t;
}
// -----------------------------------------------------------------------
bool MPC_Formulation::get_nlp_info(Index& n, Index& m, Index& nnz_jac_g,
								   Index& nnz_h_lag, IndexStyleEnum& index_style)
{
	// Number of variables
	n = (5+mode*mode-1)*P;

	// Number of constraints
	m = (5+mode*mode-1)*P;

	// number of nonzeros in the jacobian of constraints
	nnz_jac_g = (mode*mode-1)*(6*P-1)+8*P-1;

	// Number of nonzeros in the hessian of the lagrangian
	nnz_h_lag = 4;

	// We use the standard fortran index style for row/col entries
	index_style = C_STYLE;

	return true;
}
// -----------------------------------------------------------------------
bool MPC_Formulation::get_bounds_info(Index n, Number* x_l, Number* x_u,
									  Index m, Number* g_l, Number* g_u)
{
	// here, the n and m we gave IPOPT in get_nlp_info are passed back to us.
	// If desired, we could assert to make sure they are what we think they are.
	// assert(n == (5+mode*(mode+1)/2-1)*P);
	// assert(m == (5+mode*(mode+1)/2-1)*P);

	int i;

	// Bounds on x
	for(i=0;i<P;i++){
		x_l[i] = Tl;
		x_u[i] = Th;
	}
	for(i=P;i<n;i++){
		x_l[i] = 0.0;
		x_u[i] = +1e19;
	}

	// Bounds on constraints
	// Equality constraints of r_i, m_i, C2_i, sigma2_i and cov_{m,n,i}
	for(i=0;i<iCon_dT_begin;i++){
		g_l[i] = 0.0;
		g_u[i] = 0.0;
	}
	// Inequality constraints of dT
	for(i=0;i<P;i++){
		g_l[iCon_dT(i)] = dTl;
		g_u[iCon_dT(i)] = dTh;
	}

	return true;
}
// -----------------------------------------------------------------------
bool MPC_Formulation::get_starting_point(Index n, bool init_x, Number* x,
										 bool init_z, Number* z_L, Number* z_U,
										 Index m, bool init_lambda,
										 Number* lambda)
{
	// Here, we assume we only have starting values for x, if you code
	// your own NLP, you can provide starting values for the others if
	// you wish.
	assert(init_x == true);
	assert(init_z == false);
	assert(init_lambda == false);

	int i,i_mode,j_mode,index;
	// we initialize x in bounds, in the upper right quadrant
	for(i=0;i<P;i++)
	{
		x[i_T(i)] = x_old[i_T(i)];
		//x[i_T(i)] = 610;
		x[i_r(i)] = 0.0;
		x[i_m(i)] = 0.0;
		for(i_mode=0;i_mode<mode;i_mode++){
			for(j_mode=0;j_mode<mode;j_mode++){
				if(i_mode==j_mode && i_mode == 0){}
				else{
					index = i_cov(i,i_mode,j_mode);
					x[i_r(i)] += z[index2D(i_mode,j_mode)];
					x[i_m(i)] += Km2[i_mode][j_mode]*z[index2D(i_mode,j_mode)];
					x[i_cov(i,i_mode,j_mode)] = z[index2D(i_mode,j_mode)];
				}
			}
		}
		x[i_r(i)] /= (pi*pi);
		x[i_c2(i)] = CONST_COEFF_C2*exp(k_B*x[i_T(i)]/E_c);
		x[i_sigma2(i)] = CONST_COEFF_SIGMA2*(1+exp((a_t+k_t*W)*x[i_T(i)]-(a_v+k_v*W)));
	}
	

	return true;
}
// -----------------------------------------------------------------------
bool MPC_Formulation::eval_f(Index n, const Number* x, bool new_x, Number& obj_value)
{
	// return the value of the objective function
	// J = q_h*(h_set-h_f)^2+q_r*(r2_s-r2_f)^2
	int i;
	obj_value = 0.0;
	for(i=0;i<P;i++)
	{
		obj_value += q_m*(m2_set-x[i_m(i)])*(m2_set-x[i_m(i)])+q_r*(r2_set-x[i_r(i)])*(r2_set-x[i_r(i)]);
	}

	return true;
}
// -----------------------------------------------------------------------
bool MPC_Formulation::eval_grad_f(Index n, const Number* x, bool new_x, Number* grad_f)
{
	// return the gradient of the objective function grad_{x} f(x)
	int i;
	for(i = 0;i<P;i++)
	{
		// grad_{x1} f(x): x1 is not in the objective
		grad_f[i_T(i)] = 0.0;
		grad_f[i_r(i)] = 2*q_r*(x[i_r(i)]-r2_set);
		grad_f[i_m(i)] = 2*q_m*(x[i_m(i)]-m2_set);
		grad_f[i_c2(i)] = 0.0;
		grad_f[i_sigma2(i)] = 0.0;
	}
	for(i=i_cov_begin;i<n;i++)
	{
		grad_f[i] = 0.0;
	}

	return true;
}
// -----------------------------------------------------------------------
bool MPC_Formulation::eval_g(Index n, const Number* x, bool new_x, Index m, Number* g)
{
	// return the value of the constraints: g(x)
	int i,i_mode,j_mode;
	for(i=0;i<P;i++)
	{
		// Equality constraints of r_i and m_i
		g[iCon_r(i)] = 0.0;
		g[iCon_m(i)] = 0.0;
		for(i_mode=0;i_mode<mode;i_mode++){
			for(j_mode=0;j_mode<mode;j_mode++){
				if(i_mode==j_mode && i_mode == 0){}
				else
				{
					g[iCon_r(i)] += x[i_cov(i,i_mode,j_mode)];
					g[iCon_m(i)] += x[i_cov(i,i_mode,j_mode)]*Km2[i_mode][j_mode];
				}
			}
		}
		g[iCon_r(i)] = g[iCon_r(i)]/(pi*pi)-x[i_r(i)];
		g[iCon_m(i)] = g[iCon_m(i)] - x[i_m(i)];

		// Equality constraints of C2_i
		g[iCon_c2(i)] = CONST_COEFF_C2*exp(k_B*x[i_T(i)]/E_c)-x[i_c2(i)];

		// Equality constraints of sigma2_i
		g[iCon_sigma2(i)] = CONST_COEFF_SIGMA2*(1+exp((a_t+k_t*W)*x[i_T(i)]-(a_v+k_v*W)))-x[i_sigma2(i)];

		// Equality constraints of COV(m,n,i)
		double TwoReLambda;
		double cov_prev;
		for(i_mode=0;i_mode<mode;i_mode++){
			for(j_mode=0;j_mode<mode;j_mode++){
				if(i_mode==j_mode && i_mode==0){}
				else{
					TwoReLambda = -8*(i_mode*i_mode+j_mode*j_mode);
					if(i==0){ 
						cov_prev = z[index2D(i_mode,j_mode)];
					}
					else{
						cov_prev = x[i_cov(i-1,i_mode,j_mode)];
					}
					g[iCon_cov(i,i_mode,j_mode)] = cov_prev*exp(TwoReLambda*x[i_c2(i)]*dt)+x[i_sigma2(i)]*(exp(TwoReLambda*x[i_c2(i)]*dt)-1)/(TwoReLambda*x[i_c2(i)])-x[i_cov(i,i_mode,j_mode)];
				}
			}
		}

		// Inequality constraints of dT_i
		double T_prev;
		if(i==0){
			T_prev = T_old;
		}
		else{
			T_prev = x[i_T(i)-1];
		}
		g[iCon_dT(i)] = x[i_T(i)]-T_prev;
	}
	return true;
}
// -----------------------------------------------------------------------
bool MPC_Formulation::eval_jac_g(Index n, const Number* x, bool new_x,
								 Index m, Index nele_jac, Index* iRow, Index *jCol,
								 Number* values)
{
	int i;
	double TwoReLambda;
	if (values == NULL) {
		// return the structure of the jacobian of the constraints
		for(i = 0;i<P;i++){
			for(m=0;m<mode;m++){
				for(n=0;n<mode;n++){
					if(n==m && m==0){}
					else{
						iRow[iJacG_r_cov(i,m,n)] = iCon_r(i);  
						jCol[iJacG_r_cov(i,m,n)] = i_cov(i,m,n);
						
						iRow[iJacG_m_cov(i,m,n)] = iCon_m(i);  
						jCol[iJacG_m_cov(i,m,n)] = i_cov(i,m,n);
						
						iRow[iJacG_cov_c(i,m,n)] = iCon_cov(i,m,n);
						jCol[iJacG_cov_c(i,m,n)] = i_c2(i);
						
						iRow[iJacG_cov_sigma(i,m,n)] = iCon_cov(i,m,n);
						jCol[iJacG_cov_sigma(i,m,n)] = i_sigma2(i);
						
						if(i>0){
							iRow[iJacG_cov_cov_pre(i,m,n)] = iCon_cov(i,m,n);
							jCol[iJacG_cov_cov_pre(i,m,n)] = i_cov(i-1,m,n);
						}

						iRow[iJacG_cov_cov_cur(i,m,n)] = iCon_cov(i,m,n);
						jCol[iJacG_cov_cov_cur(i,m,n)] = i_cov(i,m,n);
					}
				}
			}
			iRow[iJacG_r_r(i)] = iCon_r(i);
			jCol[iJacG_r_r(i)] = i_r(i);
			iRow[iJacG_m_m(i)] = iCon_m(i);
			jCol[iJacG_m_m(i)] = i_m(i);
			iRow[iJacG_c_T(i)] = iCon_c2(i);
			jCol[iJacG_c_T(i)] = i_T(i);
			iRow[iJacG_c_c(i)] = iCon_c2(i);
			jCol[iJacG_c_c(i)] = i_c2(i);
			iRow[iJacG_sigma_T(i)] = iCon_sigma2(i);  
			jCol[iJacG_sigma_T(i)] = i_T(i);
			iRow[iJacG_sigma_sigma(i)] = iCon_sigma2(i);
			jCol[iJacG_sigma_sigma(i)] = i_sigma2(i);
			iRow[iJacG_dT_T_cur(i)] = iCon_dT(i);
			jCol[iJacG_dT_T_cur(i)] = i_T(i);
			
			if(i>0){
				iRow[iJacG_dT_T_pre(i)] = iCon_dT(i);
				jCol[iJacG_dT_T_pre(i)] = i_T(i-1);
			}
		}
	}
	else {
		// return the values of the jacobian of the constraints
		for(i=0;i<P;i++){
			for(m=0;m<mode;m++){
				for(n=0;n<mode;n++){
					if(n==m && m==0){}
					else{
						TwoReLambda = -8*(m*m+n*n);
						
						values[iJacG_r_cov(i,m,n)] = 1/(pi*pi);
						values[iJacG_m_cov(i,m,n)] = Km2[m][n];
						
						double cov_pre;
						if(i==0){cov_pre = z[index2D(m,n)];}
						else{cov_pre = x[i_cov(i-1,m,n)];}
						
						double temp1 = exp(TwoReLambda*x[i_c2(i)]*dt);
						double temp2 = (temp1*TwoReLambda*dt*TwoReLambda*x[i_c2(i)]-TwoReLambda*(temp1-1))/(TwoReLambda*TwoReLambda*x[i_c2(i)]*x[i_c2(i)]);
						values[iJacG_cov_c(i,m,n)] = cov_pre*temp1*TwoReLambda*dt+x[i_sigma2(i)]*temp2;

						values[iJacG_cov_sigma(i,m,n)] = (temp1-1)/(TwoReLambda*x[i_c2(i)]);

						if(i>0){
							values[iJacG_cov_cov_pre(i,m,n)] = temp1;
						}

						values[iJacG_cov_cov_cur(i,m,n)] = -1;
					}
				}
			}
			
			values[iJacG_r_r(i)] = -1.0;
			values[iJacG_m_m(i)] = -1.0;
			values[iJacG_c_T(i)] = CONST_COEFF_C2*exp(k_B*x[i_T(i)]/E_c)*k_B/E_c;
			values[iJacG_c_c(i)] = -1.0;
			values[iJacG_sigma_T(i)] = CONST_COEFF_SIGMA2*exp((a_t+k_t*W)*x[i_T(i)]-(a_v+k_v*W))*(a_t+k_t*W);
			values[iJacG_sigma_sigma(i)] = -1.0;
			values[iJacG_dT_T_cur(i)] = 1;
			if(i>0){values[iJacG_dT_T_pre(i)] = -1;}
		}
	}

	return true;
}
// -----------------------------------------------------------------------
bool MPC_Formulation::eval_h(Index n, const Number* x, bool new_x,
							 Number obj_factor, Index m, const Number* lambda,
							 bool new_lambda, Index nele_hess, Index* iRow,
							 Index* jCol, Number* values)
{
	if (values == NULL) {
		// return the structure. This is a symmetric matrix, fill the lower left
		// triangle only.
		iRow[0] = 0;    jCol[0] = 0;
		iRow[1] = 1;    jCol[1] = 0;
		iRow[2] = 1;    jCol[2] = 1;
	}
	else {
		// return the values
		// fill the objective portion
		values[0] = 0;
		values[1] = 0;
		values[2] = 0;
		// add the portion for the first constraint
		values[0] += 0;
		values[1] += 0;
		values[2] += 0;
	}
	return false; // USE Quasi-Newton method
}
// -----------------------------------------------------------------------
void MPC_Formulation::finalize_solution(SolverReturn status,
										Index n, const Number* x, const Number* z_L, const Number* z_U,
										Index m, const Number* g, const Number* lambda,
										Number obj_value,
										const IpoptData* ip_data,
										IpoptCalculatedQuantities* ip_cq)
{
	// here is where we would store the solution to variables, or write to a file, etc
	// so we could use the solution. Since the solution is displayed to the console,
	// we currently do nothing here.

	// Save the current solution, which will be used as the initial solution.
	int i;
	for(i = 0;i<n;i++)
	{
		x_old[i] = x[i];
	}

	printf("*** time = %f, T =",time);
	for(i = 0;i<P;i++){
		printf(" %f,",x[i]);
	}
	printf("\n");
}

// -----------------------------------------------------------------------
void MPC_Formulation::get_solution(double *x)
{
	x[0] = x_old[0];
	x[1] = W;
}
